package com.zzg.datastructure.huffmancode;

import com.zzg.datastructure.constant.GlobalConst;
import com.zzg.datastructure.huffmancode.entity.Node;

import java.io.*;
import java.util.*;

/**
 * @Author zhengzg
 * @Date 2022-09-29
 * @Version v1.0
 */
public class HuffmanCode {


    public static void main(String[] args) {
        /*
        String content = "i like like like java do you like a java abc";
        byte[] contentBytes = content.getBytes();
        //System.out.println(contentBytes.length);

        byte[] huffmanBytes = huffmanZip(contentBytes);

        System.out.println(Arrays.toString(huffmanBytes));
        byte[] sourceBytes = decode(huffmanCodes, huffmanBytes);
        System.out.println(new String(sourceBytes));

         */

        // 测试压缩
        //zipFile(GlobalConst.SRC_FILE,GlobalConst.DST_FILE);

        // 测试解压
        unZipFile(GlobalConst.UN_SRC_FILE,GlobalConst.UN_DST_FILE);
    }

    public static void unZipFile(String zipFile,String dstFile) {
        System.out.println("=========解压开始=========");
        // 自定义文件输入流
        InputStream is = null;
        // 定义一个对象输出流
        ObjectInputStream ois = null;
        // 自定义文件输出流
        OutputStream os = null;
        try {
            is = new FileInputStream(zipFile);
            ois = new ObjectInputStream(is);
            // 读取字节数组
            byte[] huffmanBytes = (byte[]) ois.readObject();
            // 读取赫夫曼编码表
            Map<Byte,String> huffmanCodes = (Map<Byte, String>) ois.readObject();

            // 解码
            byte[] decode = decode(huffmanCodes, huffmanBytes);
            // 将数组写入目标文件
            os = new FileOutputStream(dstFile);
            os.write(decode);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (Objects.nonNull(os)) {
                    os.close();
                }
                if (Objects.nonNull(ois)) {
                    ois.close();
                }
                if (Objects.nonNull(is)) {
                    is.close();
                }
                System.out.println("=========解压结束=========");
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 压缩文件
     * @param srcFile 源文件
     * @param dstFile 目录文件
     */
    public static void zipFile(String srcFile,String dstFile) {
        System.out.println("==========文件压缩开始========");
        // 创建输出流
        OutputStream os = null;
        ObjectOutputStream oos = null;
        // 创建文件输入流
        FileInputStream is = null;
        try {
            is = new FileInputStream(srcFile);
            // 创建一个和源文件大小一样的byte[]
            byte[] b = new byte[is.available()];
            // 读取文件
            is.read(b);
            // 直接对文件进行压缩
            byte[] huffmanBytes = huffmanZip(b);
            // 创建文件的输出流，存放着压缩文件
            os = new FileOutputStream(dstFile);
            oos = new ObjectOutputStream(os);
            // 把赫夫曼编码后的字节数组写入压缩文件
            oos.writeObject(huffmanBytes);
            // 注意要把赫夫曼编码写入压缩文件
            oos.writeObject(huffmanCodes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (Objects.nonNull(os)) {
                    os.close();
                }
                if (Objects.nonNull(oos)) {
                    oos.close();
                }
                if (Objects.nonNull(is)) {
                    is.close();
                }
                System.out.println("==========文件压缩结束========");
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }


    // 压缩字符串
    private static byte[] huffmanZip(byte[] bytes) {
        List<Node> nodes = getNodes(bytes);
        // System.out.println("nodes=" + nodes);

        Node huffmanTreeRoot = createHuffmanTree(nodes);
        // preOrder(huffmanTreeRoot);

        // 测试生成对应的赫夫曼编码
        getCodes(huffmanTreeRoot);
        // System.out.println(huffmanCodes);

        return zip(bytes, huffmanCodes);
    }


    // 1.统计每个字符出现的次数，Map形式存储
    private static List<Node> getNodes(byte[] bytes) {
        List<Node> nodes = new ArrayList<>();

        // 遍历bytes，统计每个字节出现的次数
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            // 如果存在count证明有字符，进行加操作
            if (Objects.isNull(count)) {
                counts.put(b, GlobalConst.START_NUMBER);
            } else {
                counts.put(b, ++count);
            }
        }
        // 把每一个键值对转成一个Node对象，加入到nodes中
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }

        return nodes;
    }

    // 2.通过List，创建对应的赫夫曼树
    private static Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            // 排序 从小到大
            Collections.sort(nodes);

            // 取出根节点权值最小的两颗二叉树
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);

            // 构建一颗新的二叉树
            Node parent = new Node(null, leftNode.getWeight() + rightNode.getWeight());
            parent.setLeft(leftNode);
            parent.setRight(rightNode);

            // 从nodes中删除处理过的二叉树
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            // 将parent加入到nodes
            nodes.add(parent);
        }

        // 返回哈夫曼树的root结点
        return nodes.get(0);
    }

    // 3.前序遍历方法
    private static void preOrder(Node root) {
        if (Objects.nonNull(root)) {
            root.preOrder();
        } else {
            System.out.println("空树，无法遍历~");
        }
    }

    // 生成赫夫曼树对应的赫夫曼编码
    // 4.将赫夫曼编码表存放在Map<Byte,String> 形式
    static Map<Byte, String> huffmanCodes = new HashMap<>();
    // 5.在生成赫夫曼编码表，需要去拼接一个路径，定义StringBuilder存储某个叶子节点的路径
    static StringBuilder stringBuilder = new StringBuilder();

    private static Map<Byte, String> getCodes(Node root) {
        if (Objects.isNull(root)) {
            return null;
        }
        // 处理root的左子树
        getCodes(root.getLeft(), GlobalConst.NO_STR, stringBuilder);
        // 处理root的右子树
        getCodes(root.getRight(), GlobalConst.YES_STR, stringBuilder);
        return huffmanCodes;
    }

    /**
     * 将传入的node节点的所有叶子节点的赫夫曼编码得到，加入到huffmanCodes集合
     *
     * @param node          传入节点
     * @param code          路径：左子节点是0，右子节点是1
     * @param stringBuilder 用于拼接路径
     */
    private static void getCodes(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder copyStringBuilder = new StringBuilder(stringBuilder);
        // 将code加入到copyStringBuilder
        copyStringBuilder.append(code);
        if (Objects.isNull(node.getData())) {
            // 非叶子节点
            // 向左递归
            getCodes(node.getLeft(), GlobalConst.NO_STR, copyStringBuilder);
            // 向右递归
            getCodes(node.getRight(), GlobalConst.YES_STR, copyStringBuilder);
        } else {
            // 叶子节点
            huffmanCodes.put(node.getData(), copyStringBuilder.toString());
        }
    }

    // 6.通过生成的赫夫曼编码表，返回一个赫夫曼编码压缩后的byte[]
    /**
     * 压缩
     * @param bytes 原始字符串对应的byte[]
     * @param huffmanCodes huffmanCodes 生成的赫夫曼map
     * @return 赫夫曼编码处理后的byte[]
     */
    private static byte[] zip(byte[] bytes,Map<Byte,String> huffmanCodes) {
        // 利用huffmanCodes将bytes转成赫夫曼编码对应的字符串
        StringBuilder stringBuilder = new StringBuilder();
        // 遍历byte数组
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodes.get(b));
        }

        int len = (stringBuilder.length() + 7) / 8;

        // 创建压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;
        for (int i = 0; i < stringBuilder.length(); i += 8) {
            String strByte;
            if (i + 8 > stringBuilder.length()) {
                strByte = stringBuilder.substring(i);
            } else {
                strByte = stringBuilder.substring(i,i + 8);
            }
            // 将strByte 转成一个byte，放入到huffmanCodeBytes
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte,2);
            index++;
        }

        return huffmanCodeBytes;
    }

    // 7.获得原始的二进制字符串
    private static String byteToBitString(boolean flag , byte b) {
        int temp = b;
        // 如果是正数我们还存在补高位
        if (flag) {
            /*
             * 这里是按位或不是视频中的与，规则是有1则一，
             * 由于一个字节取值的区间是-128~127，一个byte两个字节。（127+1） * 2正好进位，所以256与正数按位或可以补高位
             */
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        // 如果最后一位是负数，则int类型为4个字节32位，会多出24位，需要进行截取
        if (flag || temp < 0) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    // 8.对压缩数据进行解码

    /**
     * 数据解压
     *
     * @param huffmanCodes 赫夫曼编码表
     * @param huffmanBytes 赫夫曼编码得到的字节数组
     * @return 原来字符串对应的数组
     */
    private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        // 先得到huffmanBytes对应的二进制字符串
        StringBuilder stringBuilder = new StringBuilder();
        // 将byte数组转成二进制字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            byte b = huffmanBytes[i];
            // 判断是不是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBitString(!flag,b));
        }

        // 把字符串按照指定的赫夫曼编码进行解码
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(),entry.getKey());
        }

        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length();) {
            int count = 1;
            boolean flag = true;
            Byte b = null;

            while (flag) {
                // i不动，让count移动，知道匹配一个字符
                String key = stringBuilder.substring(i,i+count);
                b = map.get(key);
                if (Objects.isNull(b)) {
                    // 没有匹配到继续移动指针
                    count++;
                } else {
                    flag = false;
                }
            }

            list.add(b);
            // 让i位置移动
            i += count;
        }
        byte[] b = new byte[list.size()];
        for (int i = 0; i < b.length; i++) {
            b[i] = list.get(i);
        }
        return b;
    }

}
